import numpy as np
import pandas as pd
import os
import glob
import re
import scipy.io
import argparse
from tqdm import tqdm

def process_spatial_correlation_data(input_dir, output_dir, window_size=20, overlap=10):
    """
    处理空间相关性数据，将连续时间点的数据分割成有重叠的窗口
    
    参数:
        input_dir: 输入目录，包含空间相关性的MAT文件
        output_dir: 输出目录
        window_size: 窗口大小(每个输入的时间点数量)
        overlap: 相邻窗口的重叠时间点数量
    """
    print(f"读取空间相关性数据: {input_dir}")
    
    # 获取所有MAT文件
    spatial_files = glob.glob(os.path.join(input_dir, 'spatial_correlation_*.mat'))
    
    if not spatial_files:
        raise FileNotFoundError(f"在 {input_dir} 中没有找到空间相关性数据文件")
    
    print(f"找到 {len(spatial_files)} 个空间相关性文件")
    
    # 提取时间点ID并排序
    time_ids = []
    for file_path in spatial_files:
        match = re.search(r'spatial_correlation_(\d+)\.mat', os.path.basename(file_path))
        if match:
            time_ids.append(int(match.group(1)))
    
    time_ids = sorted(time_ids)
    print(f"时间点ID范围: {min(time_ids)} - {max(time_ids)}")
    
    # 创建一个数据字典来存储所有数据
    spatial_corr_dict = {}
    
    # 加载所有空间相关性数据
    print("加载空间相关性数据...")
    for file_path in tqdm(spatial_files):
        match = re.search(r'spatial_correlation_(\d+)\.mat', os.path.basename(file_path))
        if match:
            time_id = int(match.group(1))
            try:
                # 加载.mat文件
                mat_data = scipy.io.loadmat(file_path)
                if "p" in mat_data:
                    # 获取p值，形状通常为(1, 64)，是复数向量
                    spatial_corr = mat_data['p']
                    # 存储为字典
                    spatial_corr_dict[time_id] = spatial_corr.flatten()
            except Exception as e:
                print(f"加载文件 {file_path} 失败: {e}")
    
    print(f"成功加载 {len(spatial_corr_dict)} 个空间相关性数据点")
    
    # 确保时间点是连续的
    all_time_ids = sorted(spatial_corr_dict.keys())
    print(f"数据集中的时间点数量: {len(all_time_ids)}")
    
    # 检查gap
    gaps = []
    for i in range(1, len(all_time_ids)):
        if all_time_ids[i] - all_time_ids[i-1] > 1:
            gaps.append((all_time_ids[i-1], all_time_ids[i]))
    
    if gaps:
        print(f"警告: 发现 {len(gaps)} 个时间点间隙:")
        for start, end in gaps[:5]:  # 只显示前5个
            print(f"  间隙: {start} - {end}")
        print("这可能导致数据不连续。建议检查数据集的完整性。")
    
    # 创建数据数组
    # 每个空间相关性向量是64个复数值
    # 我们将它们拆分为实部和虚部，得到128个实数值
    # 形状: [n_samples, sequence_length, 128]
    
    # 计算窗口数量
    n_windows = (len(all_time_ids) - window_size) // (window_size - overlap) + 1
    print(f"使用窗口大小: {window_size}, 重叠: {overlap}, 可创建窗口数: {n_windows}")
    
    if n_windows < 1:
        raise ValueError(f"时间点不足以创建窗口。需要至少 {window_size} 个时间点。")
    
    # 创建空数组
    spatial_corr_real = np.zeros((n_windows, window_size, 64), dtype=np.float32)
    spatial_corr_imag = np.zeros((n_windows, window_size, 64), dtype=np.float32)
    
    # 记录窗口的时间点ID
    window_time_ids = np.zeros((n_windows, window_size), dtype=np.int32)
    
    # 填充数据
    window_idx = 0
    step = window_size - overlap
    
    for start_idx in range(0, len(all_time_ids) - window_size + 1, step):
        if window_idx >= n_windows:
            break
            
        # 获取当前窗口的时间点
        window_time_slice = all_time_ids[start_idx:start_idx + window_size]
        
        # 存储时间点ID
        window_time_ids[window_idx] = window_time_slice
        
        # 填充数据
        for t_idx, time_id in enumerate(window_time_slice):
            if time_id in spatial_corr_dict:
                complex_data = spatial_corr_dict[time_id]
                spatial_corr_real[window_idx, t_idx, :] = np.real(complex_data)
                spatial_corr_imag[window_idx, t_idx, :] = np.imag(complex_data)
            else:
                print(f"警告: 时间点 {time_id} 没有空间相关性数据，使用零填充")
                # 已经初始化为零，不需要额外操作
        
        window_idx += 1
    
    print(f"实际创建了 {window_idx} 个窗口")
    
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 保存为NPZ格式
    output_path = os.path.join(output_dir, 'spatial_correlation_data.npz')
    print(f"保存数据到: {output_path}")
    
    # 保存数据
    np.savez(
        output_path,
        spatial_corr_real=spatial_corr_real[:window_idx],  # 只保存有效窗口
        spatial_corr_imag=spatial_corr_imag[:window_idx],
        window_time_ids=window_time_ids[:window_idx]
    )
    
    print("数据处理完成!")
    print(f"实部形状: {spatial_corr_real[:window_idx].shape}")
    print(f"虚部形状: {spatial_corr_imag[:window_idx].shape}")
    
    return output_path

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='处理空间相关性数据')
    parser.add_argument('--input_dir', type=str, default='/home/zkh/lzt/damoxing/boshuyuce/xindaoyuce/dybu', help='输入目录，包含MAT文件')
    parser.add_argument('--output_dir', type=str, default='/home/zkh/lzt/damoxing/boshuyuce/xindaoyuce/dataset/spatial_corr', help='输出目录')
    parser.add_argument('--window_size', type=int, default=20, help='窗口大小(每个输入的时间点数量)')
    parser.add_argument('--overlap', type=int, default=0, help='相邻窗口的重叠时间点数量')
    
    args = parser.parse_args()
    
    process_spatial_correlation_data(
        args.input_dir,
        args.output_dir,
        args.window_size,
        args.overlap
    ) 